//! 特征模块
//! 
//! 定义市场特征和特征提取功能

use std::collections::HashMap;
use serde::{Deserialize, Serialize};
use chrono::{DateTime, Utc};
use crate::objects::bar::RawBar;

/// 特征类型枚举
#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
pub enum FeatureType {
    /// 价格特征
    Price,
    /// 成交量特征
    Volume,
    /// 技术指标特征
    Technical,
    /// 基本面特征
    Fundamental,
    /// 情绪特征
    Sentiment,
    /// 其他特征
    Other,
}

/// 特征结构体
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct Feature {
    /// 特征名称
    pub name: String,
    /// 特征类型
    pub feature_type: FeatureType,
    /// 特征值
    pub value: f64,
    /// 特征时间
    pub dt: DateTime<Utc>,
    /// 特征描述
    pub description: String,
    /// 特征数据
    pub data: HashMap<String, serde_json::Value>,
    /// 相关标的
    pub symbol: String,
    /// 特征权重
    pub weight: f64,
}

impl Feature {
    /// 创建新特征
    pub fn new(
        name: String,
        feature_type: FeatureType,
        value: f64,
        dt: DateTime<Utc>,
        description: String,
        symbol: String,
    ) -> Self {
        Self {
            name,
            feature_type,
            value,
            dt,
            description,
            data: HashMap::new(),
            symbol,
            weight: 1.0,
        }
    }
    
    /// 添加特征数据
    pub fn add_data(&mut self, key: String, value: serde_json::Value) {
        self.data.insert(key, value);
    }
    
    /// 设置特征权重
    pub fn set_weight(&mut self, weight: f64) {
        self.weight = weight;
    }
    
    /// 获取特征摘要
    pub fn get_summary(&self) -> String {
        format!(
            "{} - {} - {:.4}",
            self.name,
            self.feature_type.as_str(),
            self.value
        )
    }
    
    /// 检查特征是否有效
    pub fn is_valid(&self) -> bool {
        self.value.is_finite() && self.weight > 0.0
    }
}

impl FeatureType {
    /// 获取特征类型字符串表示
    pub fn as_str(&self) -> &'static str {
        match self {
            FeatureType::Price => "价格特征",
            FeatureType::Volume => "成交量特征",
            FeatureType::Technical => "技术指标特征",
            FeatureType::Fundamental => "基本面特征",
            FeatureType::Sentiment => "情绪特征",
            FeatureType::Other => "其他特征",
        }
    }
}

/// 特征提取器特征
pub trait FeatureExtractor {
    /// 提取特征
    fn extract_features(&self, bars: &[RawBar]) -> Vec<Feature>;
    
    /// 获取提取器名称
    fn get_name(&self) -> &str;
    
    /// 检查是否支持该特征类型
    fn supports_feature_type(&self, feature_type: FeatureType) -> bool;
}

/// 价格特征提取器
pub struct PriceFeatureExtractor {
    /// 提取器配置
    config: HashMap<String, serde_json::Value>,
}

impl PriceFeatureExtractor {
    /// 创建新的价格特征提取器
    pub fn new(config: HashMap<String, serde_json::Value>) -> Self {
        Self { config }
    }
    
    /// 提取价格特征
    fn extract_price_features(&self, bars: &[RawBar]) -> Vec<Feature> {
        if bars.is_empty() {
            return vec![];
        }
        
        let mut features = Vec::new();
        let current_bar = &bars[bars.len() - 1];
        
        // 计算价格变化率
        if bars.len() > 1 {
            let prev_bar = &bars[bars.len() - 2];
            let price_change = (current_bar.close - prev_bar.close) / prev_bar.close;
            
            let mut feature = Feature::new(
                "价格变化率".to_string(),
                FeatureType::Price,
                price_change,
                current_bar.dt,
                format!("价格变化率: {:.4}", price_change),
                current_bar.symbol.clone(),
            );
            
            feature.add_data("price_change".to_string(), serde_json::Value::Number(serde_json::Number::from_f64(price_change).unwrap()));
            features.push(feature);
        }
        
        // 计算价格波动率
        if bars.len() >= 20 {
            let recent_bars = &bars[bars.len() - 20..];
            let returns: Vec<f64> = recent_bars.windows(2)
                .map(|w| (w[1].close - w[0].close) / w[0].close)
                .collect();
            
            let mean_return = returns.iter().sum::<f64>() / returns.len() as f64;
            let variance = returns.iter()
                .map(|r| (r - mean_return).powi(2))
                .sum::<f64>() / returns.len() as f64;
            let volatility = variance.sqrt();
            
            let mut feature = Feature::new(
                "价格波动率".to_string(),
                FeatureType::Price,
                volatility,
                current_bar.dt,
                format!("价格波动率: {:.4}", volatility),
                current_bar.symbol.clone(),
            );
            
            feature.add_data("volatility".to_string(), serde_json::Value::Number(serde_json::Number::from_f64(volatility).unwrap()));
            features.push(feature);
        }
        
        // 计算价格位置
        if bars.len() >= 20 {
            let recent_bars = &bars[bars.len() - 20..];
            let high = recent_bars.iter().map(|b| b.high).fold(f64::NEG_INFINITY, f64::max);
            let low = recent_bars.iter().map(|b| b.low).fold(f64::INFINITY, f64::min);
            
            if high != low {
                let price_position = (current_bar.close - low) / (high - low);
                
                let mut feature = Feature::new(
                    "价格位置".to_string(),
                    FeatureType::Price,
                    price_position,
                    current_bar.dt,
                    format!("价格位置: {:.4}", price_position),
                    current_bar.symbol.clone(),
                );
                
                feature.add_data("price_position".to_string(), serde_json::Value::Number(serde_json::Number::from_f64(price_position).unwrap()));
                features.push(feature);
            }
        }
        
        features
    }
}

impl FeatureExtractor for PriceFeatureExtractor {
    fn extract_features(&self, bars: &[RawBar]) -> Vec<Feature> {
        self.extract_price_features(bars)
    }
    
    fn get_name(&self) -> &str {
        "PriceFeatureExtractor"
    }
    
    fn supports_feature_type(&self, feature_type: FeatureType) -> bool {
        feature_type == FeatureType::Price
    }
}

/// 成交量特征提取器
pub struct VolumeFeatureExtractor {
    /// 提取器配置
    config: HashMap<String, serde_json::Value>,
}

impl VolumeFeatureExtractor {
    /// 创建新的成交量特征提取器
    pub fn new(config: HashMap<String, serde_json::Value>) -> Self {
        Self { config }
    }
    
    /// 提取成交量特征
    fn extract_volume_features(&self, bars: &[RawBar]) -> Vec<Feature> {
        if bars.is_empty() {
            return vec![];
        }
        
        let mut features = Vec::new();
        let current_bar = &bars[bars.len() - 1];
        
        // 计算成交量变化率
        if bars.len() > 1 {
            let prev_bar = &bars[bars.len() - 2];
            let volume_change = (current_bar.vol - prev_bar.vol) / prev_bar.vol;
            
            let mut feature = Feature::new(
                "成交量变化率".to_string(),
                FeatureType::Volume,
                volume_change,
                current_bar.dt,
                format!("成交量变化率: {:.4}", volume_change),
                current_bar.symbol.clone(),
            );
            
            feature.add_data("volume_change".to_string(), serde_json::Value::Number(serde_json::Number::from_f64(volume_change).unwrap()));
            features.push(feature);
        }
        
        // 计算成交量比率
        if bars.len() >= 20 {
            let recent_bars = &bars[bars.len() - 20..];
            let avg_volume = recent_bars.iter().map(|b| b.vol).sum::<f64>() / recent_bars.len() as f64;
            let volume_ratio = current_bar.vol / avg_volume;
            
            let mut feature = Feature::new(
                "成交量比率".to_string(),
                FeatureType::Volume,
                volume_ratio,
                current_bar.dt,
                format!("成交量比率: {:.4}", volume_ratio),
                current_bar.symbol.clone(),
            );
            
            feature.add_data("volume_ratio".to_string(), serde_json::Value::Number(serde_json::Number::from_f64(volume_ratio).unwrap()));
            features.push(feature);
        }
        
        features
    }
}

impl FeatureExtractor for VolumeFeatureExtractor {
    fn extract_features(&self, bars: &[RawBar]) -> Vec<Feature> {
        self.extract_volume_features(bars)
    }
    
    fn get_name(&self) -> &str {
        "VolumeFeatureExtractor"
    }
    
    fn supports_feature_type(&self, feature_type: FeatureType) -> bool {
        feature_type == FeatureType::Volume
    }
}

/// 特征管理器
pub struct FeatureManager {
    /// 特征提取器列表
    extractors: Vec<Box<dyn FeatureExtractor>>,
    /// 特征历史
    features: Vec<Feature>,
    /// 配置
    config: HashMap<String, serde_json::Value>,
}

impl FeatureManager {
    /// 创建新的特征管理器
    pub fn new(config: HashMap<String, serde_json::Value>) -> Self {
        Self {
            extractors: Vec::new(),
            features: Vec::new(),
            config,
        }
    }
    
    /// 添加特征提取器
    pub fn add_extractor(&mut self, extractor: Box<dyn FeatureExtractor>) {
        self.extractors.push(extractor);
    }
    
    /// 提取特征
    pub fn extract_features(&mut self, bars: &[RawBar]) -> Vec<Feature> {
        let mut all_features = Vec::new();
        
        for extractor in &self.extractors {
            let features = extractor.extract_features(bars);
            all_features.extend(features);
        }
        
        // 添加到特征历史
        self.features.extend(all_features.clone());
        
        all_features
    }
    
    /// 获取特征历史
    pub fn get_features(&self) -> &[Feature] {
        &self.features
    }
    
    /// 获取特定类型的特征
    pub fn get_features_by_type(&self, feature_type: FeatureType) -> Vec<&Feature> {
        self.features.iter().filter(|f| f.feature_type == feature_type).collect()
    }
    
    /// 清理旧特征
    pub fn cleanup_old_features(&mut self, days: i64) {
        let cutoff = Utc::now() - chrono::Duration::days(days);
        self.features.retain(|f| f.dt > cutoff);
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    
    #[test]
    fn test_feature_creation() {
        let feature = Feature::new(
            "测试特征".to_string(),
            FeatureType::Price,
            0.05,
            Utc::now(),
            "测试特征描述".to_string(),
            "000001.SZ".to_string(),
        );
        
        assert_eq!(feature.name, "测试特征");
        assert_eq!(feature.feature_type, FeatureType::Price);
        assert_eq!(feature.value, 0.05);
        assert!(feature.is_valid());
    }
    
    #[test]
    fn test_price_feature_extractor() {
        let config = HashMap::new();
        let extractor = PriceFeatureExtractor::new(config);
        
        // 创建测试K线数据
        let mut bars = Vec::new();
        for i in 0..25 {
            let bar = RawBar {
                symbol: "000001.SZ".to_string(),
                dt: Utc::now() + chrono::Duration::days(i),
                open: 10.0,
                close: 10.0 + i as f64 * 0.1,
                high: 11.0,
                low: 9.0,
                vol: 1000.0,
                amount: 10000.0,
                cache: None,
            };
            bars.push(bar);
        }
        
        let features = extractor.extract_features(&bars);
        assert!(!features.is_empty());
        
        // 检查是否有价格特征
        let price_features: Vec<_> = features.iter()
            .filter(|f| f.feature_type == FeatureType::Price)
            .collect();
        assert!(!price_features.is_empty());
    }
    
    #[test]
    fn test_volume_feature_extractor() {
        let config = HashMap::new();
        let extractor = VolumeFeatureExtractor::new(config);
        
        // 创建测试K线数据
        let mut bars = Vec::new();
        for i in 0..25 {
            let bar = RawBar {
                symbol: "000001.SZ".to_string(),
                dt: Utc::now() + chrono::Duration::days(i),
                open: 10.0,
                close: 10.0,
                high: 11.0,
                low: 9.0,
                vol: 1000.0 + i as f64 * 100.0,
                amount: 10000.0,
                cache: None,
            };
            bars.push(bar);
        }
        
        let features = extractor.extract_features(&bars);
        assert!(!features.is_empty());
        
        // 检查是否有成交量特征
        let volume_features: Vec<_> = features.iter()
            .filter(|f| f.feature_type == FeatureType::Volume)
            .collect();
        assert!(!volume_features.is_empty());
    }
    
    #[test]
    fn test_feature_manager() {
        let mut manager = FeatureManager::new(HashMap::new());
        
        // 添加特征提取器
        let price_extractor = Box::new(PriceFeatureExtractor::new(HashMap::new()));
        let volume_extractor = Box::new(VolumeFeatureExtractor::new(HashMap::new()));
        
        manager.add_extractor(price_extractor);
        manager.add_extractor(volume_extractor);
        
        // 创建测试K线数据
        let mut bars = Vec::new();
        for i in 0..25 {
            let bar = RawBar {
                symbol: "000001.SZ".to_string(),
                dt: Utc::now() + chrono::Duration::days(i),
                open: 10.0,
                close: 10.0 + i as f64 * 0.1,
                high: 11.0,
                low: 9.0,
                vol: 1000.0 + i as f64 * 100.0,
                amount: 10000.0,
                cache: None,
            };
            bars.push(bar);
        }
        
        let features = manager.extract_features(&bars);
        assert!(!features.is_empty());
        
        // 检查特征历史
        assert_eq!(manager.get_features().len(), features.len());
    }
} 